"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2024-05-20
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import multiprocessing
import signal
import sys
import json

sys.path.append("../../../")

from ctypes import *
import argparse
import common.python.common as common
import struct
import pylynchipsdk as sdk
from common.python.infer_process import *
from common.python.callback_data_struct import *
import ctypes
import byte_tracker

# Create a shared boolean value to stop running thread
cancel_flag = multiprocessing.Value("b", False)
pythonapi.PyCapsule_GetPointer.restype = c_void_p
pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]

BOXNUM = 128

class bbox(ctypes.Structure):
    _fields_ = [
        ("xmin", ctypes.c_uint32),
        ("ymin", ctypes.c_uint32),
        ("xmax", ctypes.c_uint32),
        ("ymax", ctypes.c_uint32),
        ("score", ctypes.c_float),
        ("id", ctypes.c_int),
        ("label", ctypes.c_char * 64),
    ]


class Box(ctypes.Structure):
    _fields_ = [("boxesnum", ctypes.c_uint16), ("boxes", ctypes.ARRAY(bbox, BOXNUM))]


class LabelInfo(ctypes.Structure):
    _fields_ = [("name", ctypes.c_char * 64)]


class LabelList(ctypes.Structure):
    _fields_ = [("labels", ctypes.c_void_p), ("labelNum", ctypes.c_int)]


class YoloxPostProcessInfo_t(ctypes.Structure):
    _fields_ = [
        ("height", ctypes.c_int),
        ("width", ctypes.c_int),
        ("ori_height", ctypes.c_int),
        ("ori_width", ctypes.c_int),
        ("score_threshold", ctypes.c_float),
        ("nms_threshold", ctypes.c_float),
        ("nms_top_k", ctypes.c_int),
        ("is_pad_resize", ctypes.c_int),
        ("output_tensor", ctypes.c_void_p),
        ("class_num", ctypes.c_int),
        ("labelList", LabelList),
        ("boxesInfo", ctypes.c_void_p),
        ("anchorSize", ctypes.c_int),
    ]


class TrackerArgs:
    def __init__(self) -> None:
        self.track_thresh = 0.5
        self.high_thresh = 0.6
        self.match_thresh = 0.8
        self.mot20 = False
        self.track_buffer = 30


def track_callback(args):
    tracker = args[0]
    box_info_dev = args[1]
    box_size = ctypes.sizeof(Box)
    new_box_dev = args[2]
    codec_para = args[3]

    box_info_host_np = np.zeros(box_size, dtype=np.uint8)
    box_info_host_ptr = sdk.lyn_numpy_to_ptr(box_info_host_np)

    common.error_check(
        sdk.lyn_memcpy(
            box_info_host_ptr,
            box_info_dev,
            box_size,
            sdk.lyn_memcpy_dir_t.ServerToClient,
        )
    )

    box_info = ctypes.cast(
        pythonapi.PyCapsule_GetPointer(box_info_host_ptr, None), ctypes.POINTER(Box)
    ).contents

    box_array = np.array(
        [
            (
                box_info.boxes[i].xmin,
                box_info.boxes[i].ymin,
                box_info.boxes[i].xmax,
                box_info.boxes[i].ymax,
                box_info.boxes[i].score,
            )
            for i in range(box_info.boxesnum)
        ],
        dtype=np.float64,
    )

    output = tracker.update(box_array, (1, 1), (1, 1))
    boxes_info_index = 0

    new_box = Box()
    for i in output:
        vertical = i.tlwh[2] / i.tlwh[3] > 1.6
        if i.tlwh[2] * i.tlwh[3] > 20 and not vertical:
            if boxes_info_index >= BOXNUM:
                break

            new_box.boxes[boxes_info_index].xmin = int(max(0.0, i.tlbr[0]))
            new_box.boxes[boxes_info_index].ymin = int(max(0.0, i.tlbr[1]))
            new_box.boxes[boxes_info_index].xmax = int(
                min(float(codec_para.width) - 1.0, i.tlbr[2])
            )
            new_box.boxes[boxes_info_index].ymax = int(
                min(float(codec_para.height) - 1.0, i.tlbr[3])
            )
            new_box.boxes[boxes_info_index].label = str(i.track_id).encode("utf-8")
            new_box.boxes[boxes_info_index].score = i.score

            boxes_info_index += 1

    # box_info.boxesnum = boxes_info_index

    new_box.boxesnum = boxes_info_index

    new_box_np = np.frombuffer(
        new_box,
        dtype=np.uint8,
        count=ctypes.sizeof(Box) // np.dtype(np.uint8).itemsize,
    )

    new_box_ptr = sdk.lyn_numpy_to_ptr(new_box_np)
    common.error_check(
        sdk.lyn_memcpy(
            new_box_dev,
            new_box_ptr,
            box_size,
            sdk.lyn_memcpy_dir_t.ClientToServer,
        )
    )

    return 0


def create_label_list_with_names(names):
    num_labels = len(names)

    label_info_array = (LabelInfo * num_labels)()

    for i, name in enumerate(names):
        label_info_array[i].name = name["desc"].encode("utf-8")

    label_info_dev, ret = sdk.lyn_malloc(ctypes.sizeof(LabelInfo) * num_labels)
    common.error_check(ret)

    label_info_array_np = np.frombuffer(
        label_info_array,
        dtype=np.uint8,
        count=ctypes.sizeof(LabelInfo) * num_labels // np.dtype(np.uint8).itemsize,
    )

    label_info_array_ptr = sdk.lyn_numpy_to_ptr(label_info_array_np)

    common.error_check(
        sdk.lyn_memcpy(
            label_info_dev,
            label_info_array_ptr,
            ctypes.sizeof(LabelInfo) * num_labels,
            sdk.lyn_memcpy_dir_t.ClientToServer,
        )
    )

    label_info_dev_ptr = pythonapi.PyCapsule_GetPointer(label_info_dev, None)

    label_list = LabelList()
    label_list.labels = label_info_dev_ptr
    # label_list.labels = ctypes.cast(label_info_array, ctypes.POINTER(LabelInfo))
    label_list.labelNum = num_labels

    return label_list, label_info_dev


def save_boxinfo_cb(params):
    output_path = params[0]
    boxex_info = params[1]
    eos = params[2]
    if eos:
        return 0
    dst_img_size = ctypes.sizeof(Box)
    host_buf_arr = np.ones(dst_img_size, dtype=np.uint8)
    host_buf = sdk.lyn_numpy_to_ptr(host_buf_arr)
    ret = sdk.lyn_memcpy(
        host_buf, boxex_info, dst_img_size, sdk.lyn_memcpy_dir_t.ServerToClient
    )
    common.error_check(ret, "save_boxinfo_cb lyn_memcpy")
    host_buf_c = pythonapi.PyCapsule_GetPointer(host_buf, None)
    c_box = ctypes.cast(host_buf_c, ctypes.POINTER(Box)).contents
    dump_box_json(output_path, c_box)
    return 0


def dump_box_json(output_path, c_box):
    box_dict = {"boxesnum": c_box.boxesnum, "boxes": []}
    for i in range(c_box.boxesnum):
        bbox_instance = c_box.boxes[i]
        bbox_dict = {
            "xmin": bbox_instance.xmin,
            "ymin": bbox_instance.ymin,
            "xmax": bbox_instance.xmax,
            "ymax": bbox_instance.ymax,
            "score": bbox_instance.score,
            "label": bbox_instance.label.decode("utf-8"),
        }
        box_dict["boxes"].append(bbox_dict)

    # Serialize the dictionary to JSON
    with open(output_path, "a", encoding="utf-8") as f:
        f.seek(0, os.SEEK_END)
        file_size = f.tell()
        if file_size < 500 * 1024 * 1024:
            json.dump(box_dict, f, indent=2, ensure_ascii=False)
            f.write("\n")
        else:
            print("only support 500M.", output_path, "now file size:", file_size)


class yolov5(infer_process):
    def __init__(self, attr):
        super().__init__(attr)
        self.class_num = self.model_desc.outputTensorAttrArray[0].dims[2] - 5
        self.anchor_size = 1
        self.osd_plugin = attr.osd_plugin
        self.post_plugin = attr.post_plugin
        self.label_file = attr.label_file

        self.boxes_info, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
        common.error_check(ret)
        self.new_box_dev, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
        common.error_check(ret, "lyn_malloc")

        self.post_plugin, ret = sdk.lyn_plugin_register(self.post_plugin)
        common.error_check(ret, "lyn_plugin_register")
        self.osd_plugin, ret = sdk.lyn_plugin_register(self.osd_plugin)
        common.error_check(ret, "lyn_plugin_register")

        with open(self.label_file) as f:
            data = json.loads(f.read())

        self.label_list, self.label_info_dev = create_label_list_with_names(data)

        tracker_args = TrackerArgs()
        self.tracker = byte_tracker.BYTETracker(tracker_args)

    def before_ctx_destroy(self):
        common.error_check(sdk.lyn_free(self.boxes_info))
        common.error_check(sdk.lyn_free(self.label_info_dev))
        common.error_check(sdk.lyn_free(self.new_box_dev))

        ret = sdk.lyn_plugin_unregister(self.post_plugin)
        common.error_check(ret, "lyn_plugin_unregister")
        ret = sdk.lyn_plugin_unregister(self.osd_plugin)
        common.error_check(ret, "lyn_plugin_unregister")

    def plugin_process(self, apu_output_data, cb_data):
        # 等待apu处理完成
        ret = sdk.lyn_record_event(self.apu_stream, self.apu_event)
        common.error_check(ret, "lyn_record_event")
        ret = sdk.lyn_stream_wait_event(self.plugin_stream, self.apu_event)
        common.error_check(ret, "lyn_stream_wait_event")
        format = int(sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)

        apu_data_ptr = pythonapi.PyCapsule_GetPointer(apu_output_data, None)
        frame_data_ptr = pythonapi.PyCapsule_GetPointer(cb_data.frame.data, None)
        boxes_info_ptr = pythonapi.PyCapsule_GetPointer(self.boxes_info, None)
        new_box_ptr = pythonapi.PyCapsule_GetPointer(self.new_box_dev, None)

        post_para = YoloxPostProcessInfo_t()
        post_para.height = self.model_height
        post_para.width = self.model_width
        post_para.ori_height = self.codec_para.height
        post_para.ori_width = self.codec_para.width
        post_para.score_threshold = 0.25
        post_para.nms_threshold = 0.45
        post_para.nms_top_k = 500
        post_para.is_pad_resize = 1
        post_para.output_tensor = apu_data_ptr
        post_para.class_num = self.label_list.labelNum
        post_para.labelList = self.label_list
        post_para.boxesInfo = boxes_info_ptr
        post_para.anchorSize = self.anchor_size

        yolox_data = ctypes.string_at(
            ctypes.addressof(post_para), ctypes.sizeof(post_para)
        )
        ret = sdk.lyn_plugin_run_async(
            self.plugin_stream,
            self.post_plugin,
            "lynYoloxPostProcess",
            yolox_data,
            ctypes.sizeof(YoloxPostProcessInfo_t),
        )

        common.error_check(ret, "lyn_plugin_run_async")

        common.error_check(
            sdk.lyn_stream_add_callback(
                self.plugin_stream,
                track_callback,
                [self.tracker, self.boxes_info, self.new_box_dev, self.codec_para],
            )
        )

        if self.attr.show_type == 2:
            ret = sdk.lyn_stream_add_callback(
                self.plugin_stream,
                save_boxinfo_cb,
                [self.attr.output_path, self.new_box_dev, cb_data.frame.eos],
            )
        else:
            osd_para = struct.pack(
                "1P3i1P4i",
                new_box_ptr,
                self.codec_para.width,
                self.codec_para.height,
                sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
                frame_data_ptr,
                24,
                4,
                2,
                4,
            )

            common.error_check(
                sdk.lyn_plugin_run_async(
                    self.plugin_stream,
                    self.osd_plugin,
                    "lynDrawBoxAndText",
                    osd_para,
                    len(osd_para),
                )
            )

        ret = sdk.lyn_stream_add_callback(
            self.plugin_stream,
            free_to_pool_callback,
            [self.apu_output_mem_pool, apu_output_data],
        )
        common.error_check(ret, "lyn_stream_add_callback")

    def run(self, attr: infer_process_attr):
        super().run(attr)


def get_channel_infos(args):
    channel_infos = []
    device_ids = [int(id) for id in args.device_id.split(",")]

    if args.config_file is not None:
        with open(args.config_file, "r") as json_file:
            config_data = json.load(json_file)
        config_list = config_data.get("config", [])
        for config_item in config_list:
            device_id = config_item.get("devicesID", None)
            channels = config_item.get("channels", [])
            channel_idx = 0
            for channel in channels:
                channel_num = channel.get("channel", 1)
                for i in range(channel_num):
                    channel_info = copy.deepcopy(args)
                    channel_info.input_path = channel.get("input", None)
                    channel_info.output_path = channel.get("output", None)
                    channel_info.channel_id = channel_idx
                    channel_idx += 1
                    channel_info.device_id = device_id
                    channel_info.model_path = channel.get("modelPath", None)
                    channel_info.another_model_path = channel.get(
                        "anotherModelPath", None
                    )
                    channel_info.osd_plugin = channel.get("osdPlugin", None)
                    channel_info.post_plugin = channel.get("postPlugin", None)
                    channel_info.label_file = channel.get("labelFile", None)
                    channel_info.show_type = config_data.get("showType", None)
                    channel_infos.append(channel_info)
    else:
        # 若不通过json配置，则channel_info相当于args
        for device_id in range(len(device_ids)):
            for channel in range(args.channel_num):
                channel_info = copy.deepcopy(args)
                channel_info.device_id = device_id
                channel_info.channel_id = channel
                channel_infos.append(channel_info)
    return channel_infos


def main(args, video_frame) -> None:
    attr = infer_process_attr()
    attr.url = args.input_path
    if args.show_type == 1:
        attr.output_path = common.append_text_to_filename(
            common.change_file_extension(args.input_path, ".264"),
            f"_device_{args.device_id}_channel_{args.channel_id}",
        )
    elif args.show_type == 2:
        attr.output_path = common.append_text_to_filename(
            common.change_file_extension(args.input_path, ".json"),
            f"_device_{args.device_id}_channel_{args.channel_id}",
        )
    if not os.path.isfile(attr.url):
        attr.output_path = attr.output_path.replace("/", "_").replace(":", "_")
    if os.path.exists(attr.output_path):
        os.remove(attr.output_path)
    attr.device_id = args.device_id
    attr.chan_id = args.channel_id
    attr.plugin_path = args.osd_plugin
    attr.model_path = args.model_path
    attr.show_type = args.show_type
    attr.video_frame = video_frame
    attr.osd_plugin = args.osd_plugin
    attr.post_plugin = args.post_plugin
    attr.label_file = args.label_file
    if not attr.url:
        raise ValueError("input file miss!!!")
    if attr.show_type == 1 and not attr.output_path:
        raise ValueError("unspecified output path!!!")
    yolov5_worker = yolov5(attr)
    yolov5_worker.run(cancel_flag)
    yolov5_worker.close()


def cancel_process(signum, frame):
    global cancel_flag
    cancel_flag.value = True


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="python3 main.py")
    argparser.add_argument(
        "-m",
        "--model_path",
        default="../../../model/yolov6m/Net_0/",
        type=str,
        help="model file path",
    )
    argparser.add_argument(
        "-i",
        "--input_path",
        default="../../../data/yolov6m_bytetrack_with_video_sample/MOT16-09.mp4",
        type=str,
        help="input video path or stream url",
    )
    argparser.add_argument(
        "-p",
        "--post_plugin",
        default="../../../common/plugin/postprocess/cross_compiled_lib/libYolov5X7PostPlugin.so",
        type=str,
        help="plugin path",
    )
    argparser.add_argument(
        "-o",
        "--osd_plugin",
        default="../../../common/plugin/osd/lib/libOsdPlugin.so",
        type=str,
        help="plugin path",
    )
    argparser.add_argument(
        "-l",
        "--label_file",
        default="../../../common/plugin/postprocess/yolox/label/label.v5x7.json",
        type=str,
        help="label file",
    )
    argparser.add_argument(
        "-d",
        "--device_id",
        type=str,
        default="0",
        help='Comma-separated list of device ids, e.g., "0,1"',
    )
    argparser.add_argument("-s", "--show_type", type=int, default=1, help="show type")
    argparser.add_argument(
        "-c", "--channel_num", type=int, default=1, help="channel_num"
    )
    argparser.add_argument(
        "-j", "--config_file", type=str, default=None, help="config path"
    )

    signal.signal(signal.SIGINT, cancel_process)
    args = argparser.parse_args()
    channel_infos = get_channel_infos(args)
    start_process(main, channel_infos, cancel_flag)
